package com.example.zookeeper.tools;


import com.google.gson.*;
import gnu.getopt.Getopt;
import org.apache.curator.RetryPolicy;
import org.apache.curator.framework.CuratorFramework;
import org.apache.curator.framework.CuratorFrameworkFactory;
import org.apache.curator.retry.RetryNTimes;

import java.io.IOException;
import java.nio.charset.StandardCharsets;
import java.nio.file.Files;
import java.nio.file.Paths;
import java.nio.file.StandardOpenOption;
import java.util.List;
import java.util.Set;


public class ZookeeperToolsApplication {

    public static void main(String[] args) {

        //后面不带:表示不带参数如a  -a
        //后面带:表示必须带参数如a: -a xxx
        //后面带::表示可选参数如a:: -a 或者 -ahello
        Getopt opt = new Getopt("ZookeeperToolsApplication", args, "hi:o::z:t::");

        int c;
        String filename = null;
        String server = null;
        String output = null;
        String topLevel = "/config";
        while ((c = opt.getopt()) != -1) {
            switch (c) {
                case 'i':
                    filename = opt.getOptarg();
                    break;
                case 'z':
                    server = opt.getOptarg();
                    break;
                case 'o':
                    output = opt.getOptarg();
                    if(output == null) {
                        output = "console://";
                    }
                    break;
                case 't':
                    topLevel = opt.getOptarg();
                    System.out.println(topLevel);
                    if(topLevel == null) {
                        topLevel = "/config";
                    }
                    System.out.println(topLevel);
                    break;
                case 'h':
                    System.out.println("-z zookeeper服务器地址\n" +
                            "-i 导入文件地址,支持文件导入,控制台导入如: file://后面接文件路径 或者 json://后面接json格式字符串\n" +
                            "-t zookeeper顶层路径,用于导出时指定需要导出的zookeeper路径,如/config/topics\n" +
                            "-o 导出文件路径,如果不指定则将导出打印到控制台,指定的话需要使用格式如file://文件路径");
                    return;
            }
        }

        if(server == null) {
            System.out.println("please specify zookeeper server, for example: -z localhost:2181");
            return;
        }

        try {

            int sleepMsBetweenRetries = 100;
            int maxRetries = 3;
            RetryPolicy retryPolicy = new RetryNTimes(maxRetries, sleepMsBetweenRetries);
            CuratorFramework client = CuratorFrameworkFactory .newClient(server, retryPolicy);

            client.start();

            client.blockUntilConnected();

            if(filename != null) {
                String content = "";
                if(filename.startsWith("file://")) {
                    filename = filename.substring(7);
                    content = new String(java.nio.file.Files.readAllBytes(Paths.get(filename)), StandardCharsets.UTF_8);
                } else if(filename.startsWith("json://")) {
                    content = filename.substring(7);
                } else {
                    System.out.println("please specify -f file://filepath.json or -f json://{json content}");
                }

                doImport(client, "", new JsonParser().parse(content));
            }

            if(output != null) {
                JsonObject element = new JsonObject();

                JsonObject to = element;

                // 添加父级路径如 /config/application
                int idx = topLevel.lastIndexOf("/");
                if(idx != -1) {
                    String s = topLevel.substring(0, idx);
                    if(!s.isEmpty()) {
                        String[] vs = s.split("/");
                        for (int i = 0; i < vs.length; i++) {
                            if(!vs[i].isEmpty()) {
                                JsonObject oi = new JsonObject();
                                to.add(vs[i], oi);
                                to = oi;
                            }
                        }
                    }
                }

                doExport(client, topLevel, to);

                Gson gson = new GsonBuilder().setPrettyPrinting().create();

                if(output.startsWith("file://")) {
                    Files.write(Paths.get(output.substring(7)), gson.toJson(element).getBytes(), StandardOpenOption.CREATE, StandardOpenOption.TRUNCATE_EXISTING, StandardOpenOption.WRITE);
                } else {
                    System.out.println(gson.toJson(element));
                }
            }

        } catch (IOException e) {
            e.printStackTrace();
        } catch (InterruptedException e) {
            e.printStackTrace();
        }
    }

    public static void doImport(CuratorFramework client, String topLevel, JsonElement element ) {
        if(element.isJsonObject()) {
            JsonObject object = element.getAsJsonObject();
            Set<String> keySet = object.keySet();

            keySet.forEach(key -> {
                if(object.get(key).isJsonObject()) {
                    doImport(client, topLevel + "/" + key, object.get(key));
                } else {
                    String path = topLevel + "/" + key;
                    String value = object.get(key).getAsString();
                    System.out.println(path + ":" + value);
                    try {
                        if(client.checkExists().forPath(path) != null) {
                            client.setData().forPath(path, value.getBytes());
                        } else {
                            client.create().creatingParentsIfNeeded().forPath(path, value.getBytes());
                        }
                    } catch (Exception e) {
                        e.printStackTrace();
                    }
                }
            });
        }
    }

    public static void doExport(CuratorFramework client, String topLevel, JsonObject element ) {
        try {

            String p = topLevel;
            int idx  = topLevel.lastIndexOf("/");
            if(idx != -1) {
                p = p.substring(idx+1);
            }
            byte[] v = client.getData().forPath(topLevel);
            if(v != null) {
                element.addProperty(p, new String(v, StandardCharsets.UTF_8));
            }

            List<String> children = client.getChildren().forPath(topLevel);

            if(children != null && !children.isEmpty()) {

                JsonObject o = null;
                if(p.isEmpty()) {
                    // 顶层为 /时需要做特殊处理
                    o = element;
                } else {
                    o = new JsonObject();
                    element.add(p, o);
                }
                for (int i = 0; i < children.size(); i++) {
                    String path = children.get(i);
                    if(topLevel.endsWith("/")) {
                        path = topLevel + path;
                    } else {
                        path = topLevel + "/" + path;
                    }
                    doExport(client, path, o);
                }
            }

        } catch (Exception e) {
            e.printStackTrace();
        }
    }

}
